from loguru import logger
import torch,torch_npu
import os,json
from SubjectDetector.Yolo5 import YOLOSubjectDetector
from Segment.SlimSAM import SlimSAM
class MultiAgent:
    def __init__(self, device="npu",experiments_path='./experimenes',conf_threshold=0.5, iou_threshold=0.4):
        # 检查npu是否可用
        if not torch_npu.is_available():
            raise Exception("NPU is not available.")
        
        # 初始化logger
        self.logger = logger
        
        # log path expeiements_path/log
        log_path = os.path.join(experiments_path, "log")
        if not os.path.exists(log_path):
            os.makedirs(log_path)
        # 设置日志文件路径
        log_file = os.path.join(log_path, "log.txt")
        # 设置日志文件路径
        logger.add(log_file, rotation="10 MB", backtrace=True)
        # 设置日志级别
        logger.level("DEBUG")
        
        
        # 检查结果路径
        if not os.path.exists(experiments_path):
            os.makedirs(experiments_path)
        self.experiments_path = experiments_path
        self.logger.info(f"Experiments path: {experiments_path}")
        
        #visualize path experiments_path/visualize
        visualize_path = os.path.join(experiments_path, "visualize")
        if not os.path.exists(visualize_path):
            os.makedirs(visualize_path)
        self.visualize_path = visualize_path
        self.logger.info(f"Visualize path: {visualize_path}")
        
        # 设置 NPU 的缓存目录
        npu_cache_path = os.path.join(experiments_path, "npu_cache")
        if not os.path.exists(npu_cache_path):
            os.makedirs(npu_cache_path)
        # 设置npu的缓存目录为experiments_path/npu_cache
        os.environ["ASCEND_KERNEL_CACHE_PATH"] = npu_cache_path
        self.logger.info(f"NPU cache path: {npu_cache_path}")
        
        self.device = device
        self.logger.info(f"Device: {self.device}")
        
        # 初始化YOLO检测器
        self.yolo_detector = YOLOSubjectDetector(
            experiments_path, logger, conf_threshold, iou_threshold, device)

        # 初始化SlimSAM模型
        self.slimsam = SlimSAM(experiments_path, logger, device)

        # 记录初始化完成
        self.logger.info("MultiAgent initialized successfully.")
    
    def __call__(self, image_path):
        # 检测主体
        result_bbox = self.yolo_detector.detect(image_path)
        # 使用SlimSAM模型分割主体
        masks = self.slimsam.preprocess_image(image_path, result_bbox)
        # 可视化结果
        visualize_path = os.path.join(self.visualize_path, os.path.basename(image_path))
        self.slimsam.visualize_mask(image_path, masks, visualize_path)
        
        return masks